import os

import numpy as np
import tensorflow as tf
from sklearn.preprocessing import StandardScaler

os.environ["CUDA_VISIBLE_DEVICES"] = '2'

npy_dir = '/Users/d/Project/emotions/data/Spectrogram_EN_Var'
tf_dir = '/Users/d/Project/emotions/data/Spectrogram_EN_Var_tfrecord'

train_sessions = ['Ses01', 'Ses02', 'Ses03', 'Ses04']

other_sessions = ['Ses05']
count = 0

limit_len = 1200


def check_is_session(filename, sessions):
    for session in sessions:
        if session in filename:
            return True
    return False


def collect_train_data():
    global count
    files = os.listdir(npy_dir)
    train_list = list()
    for filename in files:
        print(filename)
        if check_is_session(filename, train_sessions) and ('impro' in filename):
            count += 1
            print('collecting: ', count)
            filepath = os.path.join(npy_dir, filename)
            train_data = np.load(filepath)
            train_list.append(train_data)
    return train_list


def get_scaler(train_list):
    print('get scaler')
    train_datas = np.vstack(train_list)
    scaler = StandardScaler().fit(train_datas)
    return scaler


def _byte_feature(value):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))


def _int64_feature(value):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))


def judge_label(file_name):
    if "neu" in file_name:
        return 0
    elif "ang" in file_name:
        return 1
    elif "hap" in file_name:
        return 2
    elif "sad" in file_name:
        return 3
    else:
        return -1


def process_an_example(norm_np, sentence_str):
    length = norm_np.shape[0]
    label_int = judge_label(sentence_str)
    example = tf.train.Example(features=tf.train.Features(feature={
        'data': _byte_feature(norm_np.astype(np.float32).tostring()),
        'len': _int64_feature(length),
        'sentence_id': _byte_feature(str.encode(sentence_str)),
        'label': _int64_feature(label_int)
    }))
    return example.SerializeToString()


def process_train_set(scaler):
    global count
    count = 0
    files = os.listdir(npy_dir)
    record_name = 'train.tfrecords'
    record_path = os.path.join(tf_dir, record_name)
    with tf.python_io.TFRecordWriter(record_path) as writer:
        for filename in files:
            if check_is_session(filename, train_sessions) and ('impro' in filename):

                filepath = os.path.join(npy_dir, filename)
                train_data = np.load(filepath)
                if train_data.shape[0] > limit_len:
                    continue
                count += 1
                print('save train: ', count)
                record_data = scaler.transform(train_data)
                writer.write(
                    process_an_example(record_data, filename))


def process_vali_test_set(scaler):
    global count
    count = 0
    files = os.listdir(npy_dir)
    vali_record_name = 'vali.tfrecords'
    test_record_name = 'test.tfrecords'
    vali_record_path = os.path.join(tf_dir, vali_record_name)
    test_record_path = os.path.join(tf_dir, test_record_name)
    vali_writer = tf.python_io.TFRecordWriter(vali_record_path)
    test_writer = tf.python_io.TFRecordWriter(test_record_path)
    for filename in files:
        if check_is_session(filename, other_sessions) and ('impro' in filename):

            filepath = os.path.join(npy_dir, filename)
            train_data = np.load(filepath)
            if train_data.shape[0] > limit_len:
                continue
            count += 1
            print('save vali or test', count)
            record_data = scaler.transform(train_data)
            if filename[-12] == 'F':
                vali_writer.write(process_an_example(record_data, filename))
            else:
                test_writer.write(process_an_example(record_data, filename))
    vali_writer.close()
    test_writer.close()


def main():
    train_list = collect_train_data()
    scaler = get_scaler(train_list)
    process_train_set(scaler)
    process_vali_test_set(scaler)


if __name__ == '__main__':
    main()
